package com.gitee.cirnochat.common.config;

import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;
import org.springframework.core.task.AsyncTaskExecutor;
import org.springframework.core.task.support.TaskExecutorAdapter;
import org.springframework.lang.NonNull;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.concurrent.ForkJoinPool;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 线程池配置
 */
@Configuration
public class ThreadPoolConfig {

    private final int CORE_POOL_SIZE = 10;

    private final int MAX_POOL_SIZE = 200;

    private final int QUEUE_CAPACITY = 1000;

    private final int KEEP_ALIVE_SECONDS = 60;

    private final String THREAD_NAME_PREFIX = "virtual-thread-";



    /**
     * 线程工厂
     * @return {@link ThreadFactory}
     */
    @Bean
    @ConditionalOnMissingBean
    public ThreadFactory threadFactory() {
        return new ThreadFactory() {
            private final AtomicInteger threadId = new AtomicInteger(0);
            private final ThreadFactory virtualFactory = Thread.ofVirtual().factory();
            @Override
            public Thread newThread(@NonNull Runnable r) {
                Thread thread = virtualFactory.newThread(r);
                thread.setName(THREAD_NAME_PREFIX + threadId.getAndIncrement());
                return thread;
            }
        };
    }

    /**
     * 线程池配置
     * @return {@link ThreadPoolTaskExecutor}
     */
    @Bean
    @Primary
    @ConditionalOnMissingBean
    public ThreadPoolTaskExecutor threadPoolTaskExecutor() {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        executor.setCorePoolSize(CORE_POOL_SIZE); // 核心线程数
        executor.setMaxPoolSize(MAX_POOL_SIZE); // 最大线程数
        executor.setQueueCapacity(QUEUE_CAPACITY); // 队列容量
        executor.setKeepAliveSeconds(KEEP_ALIVE_SECONDS); // 线程存活时间
        executor.setThreadFactory(threadFactory()); // 线程工厂
        executor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy()); // 拒绝策略
        return executor;
    }


    /**
     * 分支合并线程池
     * @return {@link ForkJoinPool}
     */
    @Bean
    @ConditionalOnMissingBean
    public ForkJoinPool forkJoinPool() {
        return new ForkJoinPool(
                Runtime.getRuntime().availableProcessors(),
                ForkJoinPool.defaultForkJoinWorkerThreadFactory,
                (t, e) -> {
                    // do nothing
                },
                true
        );
    }
    /**
     * 为了让SpringBoot在你使用 @Async 注解的时候用虚拟线程执行你的 @Async 任务。
     * @return {@link AsyncTaskExecutor}
     */
    @Bean
    @ConditionalOnMissingBean
    public AsyncTaskExecutor asyncTaskExecutor(){
        return new TaskExecutorAdapter(threadPoolTaskExecutor());
    }
}
